package logger

import (
	"bytes"
	"context"
	"github.com/gin-gonic/gin"
	"io"
	"sync/atomic"
)

type PrintLog func(ctx context.Context, al *AccessLog)

// MiddlewareBuilder 扩展gin默认打印日志
type MiddlewareBuilder struct {
	allowReqBody  *atomic.Bool
	allowRespBody *atomic.Bool
	loggerFunc    PrintLog
}

func NewBuilder(f PrintLog) *MiddlewareBuilder {
	return &MiddlewareBuilder{loggerFunc: f, allowRespBody: &atomic.Bool{},
		allowReqBody: &atomic.Bool{}}
}
func (b *MiddlewareBuilder) AllowReqBody() *MiddlewareBuilder {
	b.allowReqBody.Store(true)
	return b
}
func (b *MiddlewareBuilder) AllowRespBody() *MiddlewareBuilder {
	b.allowReqBody.Store(true)
	return b
}
func (b *MiddlewareBuilder) Build() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		url := ctx.Request.URL.String()
		if len(url) > 1024 {
			url = url[:1024]
		}
		log := AccessLog{
			Method:   ctx.Request.Method,
			RemoteIp: ctx.RemoteIP(),
			Url:      url,
		}
		if b.allowReqBody.Load() && ctx.Request.Body != nil {
			body, _ := ctx.GetRawData()
			//body 读完就没了，一定要放回去  只能操作一次
			ctx.Request.Body = io.NopCloser(bytes.NewReader(body))
			log.RespBody = string(body)
		}
		if b.allowRespBody.Load() {
			ctx.Writer = &responseWriter{
				al:             &log,
				ResponseWriter: ctx.Writer,
			}
		}

		defer func() {
			b.loggerFunc(ctx, &log)
		}()
		ctx.Next()
	}
}

type responseWriter struct {
	al *AccessLog
	gin.ResponseWriter
}

func (w *responseWriter) Write(data []byte) (int, error) {
	w.al.RespBody = string(data)
	return w.ResponseWriter.Write(data)
}
func (w *responseWriter) WriteString(str string) (int, error) {
	w.al.RespBody = str
	return w.ResponseWriter.WriteString(str)
}
func (w *responseWriter) WriteHeader(statusCode int) {
	w.al.StateCode = statusCode
	w.ResponseWriter.WriteHeader(statusCode)
}

type AccessLog struct {
	Method    string
	Url       string
	ReqBody   string
	RespBody  string
	RemoteIp  string
	StateCode int
}
